#include "DesktopMonitor.h"
#include "Utils.h"
#include <iostream>
#include <psapi.h>
#include <algorithm>
#include <cctype>

#pragma comment(lib, "user32.lib")
#pragma comment(lib, "advapi32.lib")

DesktopMonitor::DesktopMonitor() :
    m_currentDesktop(NULL),
    m_currentDesktopType(DesktopType::UNKNOWN),
    m_lastSessionId(0) {
}

DesktopMonitor::~DesktopMonitor() {
    Cleanup();
}

bool DesktopMonitor::Initialize() {
    Utils::DebugLog("DesktopMonitor: Initializing...");

    if (!Utils::EnableDebugPrivilege()) {
        Utils::DebugLog("DesktopMonitor: Warning - Could not enable debug privilege");
    }

    m_lastSessionId = Utils::GetCurrentSessionId();
    m_currentDesktop = GetThreadDesktop(GetCurrentThreadId());

    if (m_currentDesktop) {
        m_currentDesktopName = GetDesktopName();
        m_currentDesktopType = DetermineDesktopType(m_currentDesktopName);
    }

    Utils::DebugLog("DesktopMonitor: Initialized successfully");
    return true;
}

void DesktopMonitor::Cleanup() {
    Utils::DebugLog("DesktopMonitor: Cleaning up...");

    if (m_currentDesktop) {
        m_currentDesktop = NULL;
    }
}

bool DesktopMonitor::HasDesktopChanged() {
    HDESK hCurrentDesktop = GetThreadDesktop(GetCurrentThreadId());
    DWORD currentSessionId = Utils::GetCurrentSessionId();

    bool sessionChanged = (currentSessionId != m_lastSessionId);
    bool desktopChanged = (hCurrentDesktop != m_currentDesktop);

    // Also check for lock screen by trying to get foreground window
    HWND hForeground = GetForegroundWindow();
    bool possibleLockScreen = (hForeground == NULL || !IsWindowVisible(hForeground));

    if (sessionChanged || desktopChanged || possibleLockScreen) {
        std::string newDesktopName = GetDesktopName();
        DesktopType newDesktopType = DetermineDesktopType(newDesktopName);

        // Only consider it changed if desktop type actually changed
        bool typeChanged = (newDesktopType != m_currentDesktopType);

        if (sessionChanged || desktopChanged || typeChanged) {
            Utils::DebugLog("DesktopMonitor: Desktop change detected - " + newDesktopName);

            m_currentDesktop = hCurrentDesktop;
            m_lastSessionId = currentSessionId;
            m_currentDesktopName = newDesktopName;
            m_currentDesktopType = newDesktopType;

            return true;
        }
    }

    return false;
}

DesktopType DesktopMonitor::GetCurrentDesktopType() {
    return m_currentDesktopType;
}

HDESK DesktopMonitor::GetCurrentDesktop() {
    return m_currentDesktop;
}

std::string DesktopMonitor::GetDesktopName() {
    if (!m_currentDesktop) {
        return "";
    }

    char desktopName[256] = { 0 };
    DWORD nameLength = 0;

    if (!GetUserObjectInformationA(m_currentDesktop, UOI_NAME, NULL, 0, &nameLength)) {
        if (::GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
            m_lastError = "Failed to get desktop name length: " + Utils::GetLastWin32Error();
            return "";
        }
    }

    if (nameLength > sizeof(desktopName)) {
        nameLength = sizeof(desktopName) - 1;
    }

    if (!GetUserObjectInformationA(m_currentDesktop, UOI_NAME, desktopName, nameLength, &nameLength)) {
        m_lastError = "Failed to get desktop name: " + Utils::GetLastWin32Error();
        return "";
    }

    return std::string(desktopName);
}

DesktopType DesktopMonitor::DetermineDesktopType(const std::string& desktopName) {
    if (desktopName.empty()) {
        return DesktopType::UNKNOWN;
    }

    std::string lowerName = desktopName;
    std::transform(lowerName.begin(), lowerName.end(), lowerName.begin(), ::tolower);

    if (lowerName == "default") {
        // Check if we're actually on a secure desktop even if name is "default"
        HWND hwndForeground = GetForegroundWindow();
        if (!hwndForeground) {
            // No foreground window might indicate secure desktop
            return DesktopType::SECURE;
        }

        char className[256] = { 0 };
        if (GetClassNameA(hwndForeground, className, sizeof(className))) {
            std::string classNameStr(className);
            if (classNameStr == "Credential Dialog Xaml Host" ||
                classNameStr == "Windows.UI.Core.CoreWindow" ||
                classNameStr == "LogonUI Logon Window" ||
                classNameStr.find("LockScreen") != std::string::npos) {
                return DesktopType::WINLOGON;
            }
        }

        char windowText[256] = { 0 };
        if (GetWindowTextA(hwndForeground, windowText, sizeof(windowText))) {
            std::string windowTextStr(windowText);
            if (windowTextStr.find("Lock screen") != std::string::npos ||
                windowTextStr.find("Sign in") != std::string::npos ||
                windowTextStr.find("Password") != std::string::npos) {
                return DesktopType::WINLOGON;
            }
        }

        return DesktopType::NORMAL;
    } else if (lowerName == "winlogon") {
        return DesktopType::WINLOGON;
    } else if (lowerName == "secure") {
        return DesktopType::SECURE;
    } else if (lowerName.find("screensaver") != std::string::npos) {
        return DesktopType::SCREENSAVER;
    } else if (lowerName.find("uac") != std::string::npos ||
               lowerName.find("consent") != std::string::npos) {
        return DesktopType::UAC;
    }

    Utils::DebugLog("DesktopMonitor: Unknown desktop type for: " + desktopName);
    return DesktopType::UNKNOWN;
}

bool DesktopMonitor::SwitchToDesktop(HDESK hDesktop) {
    if (!hDesktop) {
        m_lastError = "Invalid desktop handle";
        return false;
    }

    if (!SetThreadDesktop(hDesktop)) {
        m_lastError = "Failed to set thread desktop: " + Utils::GetLastWin32Error();
        return false;
    }

    return true;
}